{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What is torch.nn Really?\n",
    "\n",
    "## 1. MNIST Data Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_PATH = Path('data')\n",
    "PATH = DATA_PATH / 'mnist'\n",
    "\n",
    "PATH.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "URL = 'http://deeplearning.net/data/mnist/'\n",
    "FILENAME = 'mnist.pkl.gz'\n",
    "\n",
    "if not (PATH / FILENAME).exists():\n",
    "    content = requests.get(URL + FILENAME).content\n",
    "    (PATH / FILENAME).open('wb').write(content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**[pickle -> memory]** This dataset is in numpy array format, and has been stored using pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import gzip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "with gzip.open((PATH / FILENAME).as_posix(), 'rb') as fr:\n",
    "    ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(fr, encoding='latin-1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**[Reshape]** Each image is 28 x 28, and is being stored as a flattened row of length 784 (=28x28). We need to reshape it to 2d first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(50000, 784)\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAD8CAYAAAC4nHJkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4wLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvqOYd8AAADgpJREFUeJzt3X+MVfWZx/HPs1j+kKI4aQRCYSnEYJW4082IjSWrxkzVDQZHrekkJjQapn8wiU02ZA3/VNNgyCrslmiamaZYSFpKE3VB0iw0otLGZuKIWC0srTFsO3IDNTjywx9kmGf/mEMzxbnfe+fec++5zPN+JeT+eM6558kNnznn3O+592vuLgDx/EPRDQAoBuEHgiL8QFCEHwiK8ANBEX4gKMIPBEX4gaAIPxDUZc3cmJlxOSHQYO5u1SxX157fzO40syNm9q6ZPVrPawFoLqv12n4zmybpj5I6JQ1Jel1St7sfSqzDnh9osGbs+ZdJetfd33P3c5J+IWllHa8HoInqCf88SX8Z93goe+7vmFmPmQ2a2WAd2wKQs3o+8Jvo0OJzh/Xu3i+pX+KwH2gl9ez5hyTNH/f4y5KO1dcOgGapJ/yvS7rGzL5iZtMlfVvSrnzaAtBoNR/2u/uImfVK2iNpmqQt7v6H3DoD0FA1D/XVtDHO+YGGa8pFPgAuXYQfCIrwA0ERfiAowg8ERfiBoAg/EBThB4Ii/EBQhB8IivADQRF+ICjCDwRF+IGgCD8QFOEHgiL8QFCEHwiK8ANBEX4gKMIPBEX4gaAIPxAU4QeCIvxAUIQfCIrwA0ERfiAowg8EVfMU3ZJkZkclnZZ0XtKIu3fk0RTyM23atGT9yiuvbOj2e3t7y9Yuv/zy5LpLlixJ1tesWZOsP/XUU2Vr3d3dyXU//fTTZH3Dhg3J+uOPP56st4K6wp+5zd0/yOF1ADQRh/1AUPWG3yXtNbM3zKwnj4YANEe9h/3fcPdjZna1pF+b2f+6+/7xC2R/FPjDALSYuvb87n4suz0h6QVJyyZYpt/dO/gwEGgtNYffzGaY2cwL9yV9U9I7eTUGoLHqOeyfLekFM7vwOj939//JpSsADVdz+N39PUn/lGMvU9aCBQuS9enTpyfrN998c7K+fPnysrVZs2Yl173vvvuS9SINDQ0l65s3b07Wu7q6ytZOnz6dXPett95K1l999dVk/VLAUB8QFOEHgiL8QFCEHwiK8ANBEX4gKHP35m3MrHkba6L29vZkfd++fcl6o79W26pGR0eT9YceeihZP3PmTM3bLpVKyfqHH36YrB85cqTmbTeau1s1y7HnB4Ii/EBQhB8IivADQRF+ICjCDwRF+IGgGOfPQVtbW7I+MDCQrC9atCjPdnJVqffh4eFk/bbbbitbO3fuXHLdqNc/1ItxfgBJhB8IivADQRF+ICjCDwRF+IGgCD8QVB6z9IZ38uTJZH3t2rXJ+ooVK5L1N998M1mv9BPWKQcPHkzWOzs7k/WzZ88m69dff33Z2iOPPJJcF43Fnh8IivADQRF+ICjCDwRF+IGgCD8QFOEHgqr4fX4z2yJphaQT7r40e65N0g5JCyUdlfSAu6d/6FxT9/v89briiiuS9UrTSff19ZWtPfzww8l1H3zwwWR9+/btyTpaT57f5/+ppDsveu5RSS+5+zWSXsoeA7iEVAy/u++XdPElbCslbc3ub5V0T859AWiwWs/5Z7t7SZKy26vzawlAMzT82n4z65HU0+jtAJicWvf8x81sriRltyfKLeju/e7e4e4dNW4LQAPUGv5dklZl91dJ2plPOwCapWL4zWy7pN9JWmJmQ2b2sKQNkjrN7E+SOrPHAC4hFc/53b27TOn2nHsJ69SpU3Wt/9FHH9W87urVq5P1HTt2JOujo6M1bxvF4go/ICjCDwRF+IGgCD8QFOEHgiL8QFBM0T0FzJgxo2ztxRdfTK57yy23JOt33XVXsr53795kHc3HFN0Akgg/EBThB4Ii/EBQhB8IivADQRF+ICjG+ae4xYsXJ+sHDhxI1oeHh5P1l19+OVkfHBwsW3vmmWeS6zbz/+ZUwjg/gCTCDwRF+IGgCD8QFOEHgiL8QFCEHwiKcf7gurq6kvVnn302WZ85c2bN2163bl2yvm3btmS9VCrVvO2pjHF+AEmEHwiK8ANBEX4gKMIPBEX4gaAIPxBUxXF+M9siaYWkE+6+NHvuMUmrJf01W2ydu/+q4sYY57/kLF26NFnftGlTsn777bXP5N7X15esr1+/Pll///33a972pSzPcf6fSrpzguf/093bs38Vgw+gtVQMv7vvl3SyCb0AaKJ6zvl7zez3ZrbFzK7KrSMATVFr+H8kabGkdkklSRvLLWhmPWY2aGblf8wNQNPVFH53P+7u5919VNKPJS1LLNvv7h3u3lFrkwDyV1P4zWzuuIddkt7Jpx0AzXJZpQXMbLukWyV9ycyGJH1f0q1m1i7JJR2V9N0G9gigAfg+P+oya9asZP3uu+8uW6v0WwFm6eHqffv2JeudnZ3J+lTF9/kBJBF+ICjCDwRF+IGgCD8QFOEHgmKoD4X57LPPkvXLLktfhjIyMpKs33HHHWVrr7zySnLdSxlDfQCSCD8QFOEHgiL8QFCEHwiK8ANBEX4gqIrf50dsN9xwQ7J+//33J+s33nhj2VqlcfxKDh06lKzv37+/rtef6tjzA0ERfiAowg8ERfiBoAg/EBThB4Ii/EBQjPNPcUuWLEnWe3t7k/V77703WZ8zZ86ke6rW+fPnk/VSqZSsj46O5tnOlMOeHwiK8ANBEX4gKMIPBEX4gaAIPxAU4QeCqjjOb2bzJW2TNEfSqKR+d/+hmbVJ2iFpoaSjkh5w9w8b12pclcbSu7u7y9YqjeMvXLiwlpZyMTg4mKyvX78+Wd+1a1ee7YRTzZ5/RNK/uftXJX1d0hozu07So5JecvdrJL2UPQZwiagYfncvufuB7P5pSYclzZO0UtLWbLGtku5pVJMA8jepc34zWyjpa5IGJM1295I09gdC0tV5Nwegcaq+tt/MvijpOUnfc/dTZlVNByYz65HUU1t7ABqlqj2/mX1BY8H/mbs/nz193MzmZvW5kk5MtK6797t7h7t35NEwgHxUDL+N7eJ/Iumwu28aV9olaVV2f5Wknfm3B6BRKk7RbWbLJf1G0tsaG+qTpHUaO+//paQFkv4s6VvufrLCa4Wconv27NnJ+nXXXZesP/3008n6tddeO+me8jIwMJCsP/nkk2VrO3em9xd8Jbc21U7RXfGc391/K6nci90+maYAtA6u8AOCIvxAUIQfCIrwA0ERfiAowg8ExU93V6mtra1sra+vL7lue3t7sr5o0aKaesrDa6+9lqxv3LgxWd+zZ0+y/sknn0y6JzQHe34gKMIPBEX4gaAIPxAU4QeCIvxAUIQfCCrMOP9NN92UrK9duzZZX7ZsWdnavHnzauopLx9//HHZ2ubNm5PrPvHEE8n62bNna+oJrY89PxAU4QeCIvxAUIQfCIrwA0ERfiAowg8EFWacv6urq656PQ4dOpSs7969O1kfGRlJ1lPfuR8eHk6ui7jY8wNBEX4gKMIPBEX4gaAIPxAU4QeCIvxAUObu6QXM5kvaJmmOpFFJ/e7+QzN7TNJqSX/NFl3n7r+q8FrpjQGom7tbNctVE/65kua6+wEzmynpDUn3SHpA0hl3f6rapgg/0HjVhr/iFX7uXpJUyu6fNrPDkor96RoAdZvUOb+ZLZT0NUkD2VO9ZvZ7M9tiZleVWafHzAbNbLCuTgHkquJh/98WNPuipFclrXf3581stqQPJLmkH2js1OChCq/BYT/QYLmd80uSmX1B0m5Je9x90wT1hZJ2u/vSCq9D+IEGqzb8FQ/7zcwk/UTS4fHBzz4IvKBL0juTbRJAcar5tH+5pN9IeltjQ32StE5St6R2jR32H5X03ezDwdRrsecHGizXw/68EH6g8XI77AcwNRF+ICjCDwRF+IGgCD8QFOEHgiL8QFCEHwiK8ANBEX4gKMIPBEX4gaAIPxAU4QeCavYU3R9I+r9xj7+UPdeKWrW3Vu1Lorda5dnbP1a7YFO/z/+5jZsNuntHYQ0ktGpvrdqXRG+1Kqo3DvuBoAg/EFTR4e8vePsprdpbq/Yl0VutCumt0HN+AMUpes8PoCCFhN/M7jSzI2b2rpk9WkQP5ZjZUTN728wOFj3FWDYN2gkze2fcc21m9msz+1N2O+E0aQX19piZvZ+9dwfN7F8L6m2+mb1sZofN7A9m9kj2fKHvXaKvQt63ph/2m9k0SX+U1ClpSNLrkrrd/VBTGynDzI5K6nD3wseEzexfJJ2RtO3CbEhm9h+STrr7huwP51Xu/u8t0ttjmuTMzQ3qrdzM0t9Rge9dnjNe56GIPf8ySe+6+3vufk7SLyStLKCPlufu+yWdvOjplZK2Zve3auw/T9OV6a0luHvJ3Q9k909LujCzdKHvXaKvQhQR/nmS/jLu8ZBaa8pvl7TXzN4ws56im5nA7AszI2W3Vxfcz8UqztzcTBfNLN0y710tM17nrYjwTzSbSCsNOXzD3f9Z0l2S1mSHt6jOjyQt1tg0biVJG4tsJptZ+jlJ33P3U0X2Mt4EfRXyvhUR/iFJ88c9/rKkYwX0MSF3P5bdnpD0gsZOU1rJ8QuTpGa3Jwru52/c/bi7n3f3UUk/VoHvXTaz9HOSfubuz2dPF/7eTdRXUe9bEeF/XdI1ZvYVM5su6duSdhXQx+eY2YzsgxiZ2QxJ31TrzT68S9Kq7P4qSTsL7OXvtMrMzeVmllbB712rzXhdyEU+2VDGf0maJmmLu69vehMTMLNFGtvbS2PfePx5kb2Z2XZJt2rsW1/HJX1f0n9L+qWkBZL+LOlb7t70D97K9HarJjlzc4N6Kzez9IAKfO/ynPE6l364wg+IiSv8gKAIPxAU4QeCIvxAUIQfCIrwA0ERfiAowg8E9f/Ex0YKZYOZcwAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pyplot.imshow(x_train[0].reshape((28, 28)), cmap='gray')\n",
    "print(x_train.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**[numpy array -> Tensor]** PyTorch uses torch.tensor, rather than numpy arrays, so we need to convert our data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        ...,\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.]]) tensor([5, 0, 4,  ..., 8, 4, 8])\n",
      "torch.Size([50000, 784])\n",
      "tensor(0) tensor(9)\n"
     ]
    }
   ],
   "source": [
    "x_train, y_train, x_valid, y_valid = map(torch.tensor, (x_train, y_train, x_valid, y_valid))\n",
    "n, c = x_train.shape\n",
    "print(x_train, y_train)\n",
    "print(x_train.shape)\n",
    "print(y_train.min(), y_train.max())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Neural Net From Scratch (No torch.nn)\n",
    "\n",
    "PyTorch provides methods to create random or zero-filled tensors, which we will use to create our weights and bias for a simple linear model. When setting ***requires_grad=True***, PyTorch will **record all of the operations done on the tensor**, so that it can calculate the gradient during back-propagation automatically!\n",
    "\n",
    "For the weights, we set requires_grad **after the initialization**, since we don't want that step included in the gradient!\n",
    "\n",
    "> 刘尧：设置requires_grad=True后，PyTorch**只是记录tensor上的所有操作和中间变量**，以用于随后backward时使用这些中间变量，仅此而已，没太多复杂逻辑。\n",
    "\n",
    "> 刘尧：设置之前的操作不会被包括在backward的梯度计算中，之后的操作会，除非Wrap在**with torch.no_grad()**中，所以要注意设置的时机！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = torch.randn(784, 10) / math.sqrt(784)\n",
    "weights.requires_grad_()\n",
    "bias = torch.zeros(10, requires_grad=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thanks to PyTorch's ability to calculate gradients automatically, we can use **any standard Python function (or callable object) as a model**! So let's just write a plain matrix multiplication and broadcasted addition to create a simple linear model! We also need an activation function, so we'll write *log_softmax* and use it.\n",
    "\n",
    "Remember: although PyTorch provides lots of pre-written loss functions, activation functions,and so forth, you can easily write your own **using plain Python**. PyTorch will even create fast GPU or vectorized CPU code for your function automatically!\n",
    "\n",
    "> 刘尧：PyTorch支持我们使用plain Python来创建自己的model和function！哎呀真是太方便太灵活太有爱了！我喜欢~"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_softmax(x):\n",
    "    return x - x.exp().sum(-1).log().unsqueeze(-1)  # x是Tensor 像对待numpy array那样来自定义函数！当然，api使用Tensor的\n",
    "\n",
    "def model(xb):\n",
    "    return log_softmax(xb @ weights + bias)  # @ means dot production"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-2.1582, -2.5963, -2.2912, -2.2727, -2.0091, -2.7189, -2.2418, -2.5606,\n",
      "        -2.5164, -1.9562], grad_fn=<SelectBackward>) torch.Size([64, 10])\n"
     ]
    }
   ],
   "source": [
    "bs = 64  # batch size\n",
    "xb, yb = x_train[0: bs], y_train[0: bs]\n",
    "preds = model(xb)\n",
    "print(preds[0], preds.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we see, the *preds* tensor contains not only tensor values, but also **a gradient function**, which will be used later to do backward.\n",
    "\n",
    "Let's implement negative log-likelihood to use as the loss function (again, we can just use standard Python), and then implement a function to calculate the accuracy of model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nll(input, target):\n",
    "    return - input[range(target.shape[0]), target].mean()  # target所对应的那些“列”的均值\n",
    "\n",
    "def accuracy(out, yb):\n",
    "    preds = torch.argmax(out, dim=1)\n",
    "    return (preds == yb).float().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(2.3263, grad_fn=<NegBackward>)\n",
      "tensor(0.0781)\n"
     ]
    }
   ],
   "source": [
    "loss_func = nll\n",
    "print(loss_func(preds, yb))\n",
    "print(accuracy(preds, yb))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we run a training loop. For each iteration, we will:\n",
    "\n",
    "- select a min-batch of data (of size *bs*)\n",
    "\n",
    "- use the model to make a predictions\n",
    "\n",
    "- calculate the loss\n",
    "\n",
    "- ***loss.backward()*** updates the gradients of the model, in this case, *weights* and *bias*\n",
    "\n",
    "We use these gradients to update the weights and bias within the ***torch.no_grad()*** context manager, because we don't want these actions to **be recorded for our next calculation of the gradients**!\n",
    "\n",
    "> 刘尧：设置了weight和bias的require_grad=True后，随后所有针对这两者的操作都会被record到梯度计算中，当某些操作不需要被record时，需要使用**torch.no_grad()**包裹这些操作！\n",
    "\n",
    "Before next training loop, we should set the gradients to zero. Otherwise, our gradients would record a running tally of all the operations that had happened. (i.e. *loss.backward()* **adds the gradients to whatever is already stored, rather than replacing them**).\n",
    "\n",
    "> 刘尧：loss.backward()会**累加计算**所有已存在的gradient，而非重新替换！所以**每个mini-batch(而非epoch)后**要及时对weights和bias的梯度归零！如：***weights.grad.zero_()***和***bias.grad.zero_()***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.debugger import set_trace  # "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0 --- i:    1 --- loss:  0.6862 --- acc:  0.8438 --- val_loss:  0.6312 --- val_acc:  0.8684\n",
      "Epoch: 0 --- i:  301 --- loss:  0.6626 --- acc:  0.8125 --- val_loss:  0.6216 --- val_acc:  0.8695\n",
      "Epoch: 0 --- i:  601 --- loss:  0.5753 --- acc:  0.8906 --- val_loss:  0.6125 --- val_acc:  0.8710\n",
      "Epoch: 1 --- i:    1 --- loss:  0.6610 --- acc:  0.8594 --- val_loss:  0.6073 --- val_acc:  0.8713\n",
      "Epoch: 1 --- i:  301 --- loss:  0.6440 --- acc:  0.8281 --- val_loss:  0.5991 --- val_acc:  0.8715\n",
      "Epoch: 1 --- i:  601 --- loss:  0.5545 --- acc:  0.8906 --- val_loss:  0.5913 --- val_acc:  0.8734\n",
      "Epoch: 2 --- i:    1 --- loss:  0.6392 --- acc:  0.8750 --- val_loss:  0.5867 --- val_acc:  0.8731\n",
      "Epoch: 2 --- i:  301 --- loss:  0.6281 --- acc:  0.8438 --- val_loss:  0.5797 --- val_acc:  0.8730\n",
      "Epoch: 2 --- i:  601 --- loss:  0.5365 --- acc:  0.8906 --- val_loss:  0.5729 --- val_acc:  0.8747\n",
      "Epoch: 3 --- i:    1 --- loss:  0.6200 --- acc:  0.8750 --- val_loss:  0.5689 --- val_acc:  0.8752\n",
      "Epoch: 3 --- i:  301 --- loss:  0.6142 --- acc:  0.8594 --- val_loss:  0.5628 --- val_acc:  0.8753\n",
      "Epoch: 3 --- i:  601 --- loss:  0.5206 --- acc:  0.8906 --- val_loss:  0.5569 --- val_acc:  0.8766\n",
      "Epoch: 4 --- i:    1 --- loss:  0.6030 --- acc:  0.8750 --- val_loss:  0.5534 --- val_acc:  0.8770\n",
      "Epoch: 4 --- i:  301 --- loss:  0.6020 --- acc:  0.8594 --- val_loss:  0.5479 --- val_acc:  0.8768\n",
      "Epoch: 4 --- i:  601 --- loss:  0.5065 --- acc:  0.8906 --- val_loss:  0.5427 --- val_acc:  0.8785\n"
     ]
    }
   ],
   "source": [
    "lr = 0.001\n",
    "epochs = 5\n",
    "loss_func = nll\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    for i in range((n - 1) // bs + 1):\n",
    "        set_trace\n",
    "        # Training\n",
    "        xb, yb = x_train[i * bs: (i + 1) * bs], y_train[i * bs: (i + 1) * bs]\n",
    "        pred = model(xb)\n",
    "        loss = loss_func(pred, yb)\n",
    "        acc = accuracy(pred, yb)\n",
    "\n",
    "        # Validation  注意，特别注意！这么写validation是不对的，每次train/val的loss和acc应该是每个epoch里遍历所有mini-batch后的均值！ 此处只是反例示范~啦啦啦~~\n",
    "        pred_val = model(x_valid)\n",
    "        val_loss = loss_func(pred_val, y_valid)\n",
    "        val_acc = accuracy(pred_val, y_valid)\n",
    "        \n",
    "        # Log\n",
    "        if i % 300 == 0:\n",
    "            print(f'Epoch: {epoch} --- i: {i + 1: 4d} --- loss: {loss.item(): .4f} --- acc: {acc: .4f} --- val_loss: {val_loss: .4f} --- val_acc: {val_acc: .4f}')\n",
    "        \n",
    "        loss.backward()\n",
    "        with torch.no_grad():\n",
    "            weights -= lr * weights.grad\n",
    "            bias -= lr * bias.grad\n",
    "            weights.grad.zero_()\n",
    "            bias.grad.zero_()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Using torch.nn.functional\n",
    "\n",
    "> **activation & loss functions --> torch.nn.functional**\n",
    "\n",
    "We will now refactor our code, so that it does the same things as before, only we'll start taking advantage of PyTorch's *nn* classes to make it more concise and flexible. At each step from here, we should be making our code one or more: **shorter, more understandable, and/or more flexible**.\n",
    "\n",
    "Frist and easiest step: replacing our hand-written activation and loss functions with ***torch.nn.functional***."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_func = F.cross_entropy  # change\n",
    "\n",
    "def model(xb):\n",
    "    return xb @ weights + bias  # Note that we no longer call log_softmax in the model function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Refactor using nn.Module\n",
    "\n",
    "> **model function & weights/bias variable & training loop --> nn.Moduel & nn.Parameter**\n",
    "\n",
    "We want to create a class that holds our weights, bias and method for the forward step. We will use a class instead of just using functions.\n",
    "\n",
    "> 刘尧：此处的model其实跟一般PyTorch中定义的model没本质区别，此处*\\_\\_init\\_\\_*和*forward*里是low-level的变量和操作如**weights,bias和它们加减相乘**，而一般model是high-level的比如**Layer和Lyaer(x)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mnist_Logistic(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.weights = nn.Parameter(torch.randn(784, 10) / math.sqrt(784))  # change\n",
    "        self.bias = nn.Parameter(torch.zeros(10))                           # change\n",
    "        \n",
    "    def forward(self, xb):\n",
    "        return xb @ self.weights * self.bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(2.3026, grad_fn=<NllLossBackward>)\n"
     ]
    }
   ],
   "source": [
    "model = Mnist_Logistic()\n",
    "print(loss_func(model(xb), yb))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit():\n",
    "    for epoch in range(epochs):\n",
    "        for i in range((n - 1) // bs + 1):\n",
    "            xb, yb = x_train[i * bs: (i + 1) * bs], y_train[i * bs: (i + 1) * bs]\n",
    "            pred = model(xb)\n",
    "            loss = loss_func(pred, yb)\n",
    "            \n",
    "            if i % 300 == 0:\n",
    "                print(f'Epoch: {epoch} --- i: {i + 1: 4d} --- loss: {loss.item(): .4f}')\n",
    "            \n",
    "            loss.backward()\n",
    "            with torch.no_grad():\n",
    "                for p in model.parameters():\n",
    "                    p -= lr * p.grad\n",
    "                model.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0 --- i:    1 --- loss:  2.2640\n",
      "Epoch: 0 --- i:  301 --- loss:  2.2454\n",
      "Epoch: 0 --- i:  601 --- loss:  2.2683\n",
      "Epoch: 1 --- i:    1 --- loss:  2.2420\n",
      "Epoch: 1 --- i:  301 --- loss:  2.2118\n",
      "Epoch: 1 --- i:  601 --- loss:  2.2480\n",
      "Epoch: 2 --- i:    1 --- loss:  2.2087\n",
      "Epoch: 2 --- i:  301 --- loss:  2.1616\n",
      "Epoch: 2 --- i:  601 --- loss:  2.2176\n",
      "Epoch: 3 --- i:    1 --- loss:  2.1598\n",
      "Epoch: 3 --- i:  301 --- loss:  2.0897\n",
      "Epoch: 3 --- i:  601 --- loss:  2.1730\n",
      "Epoch: 4 --- i:    1 --- loss:  2.0906\n",
      "Epoch: 4 --- i:  301 --- loss:  1.9914\n",
      "Epoch: 4 --- i:  601 --- loss:  2.1087\n"
     ]
    }
   ],
   "source": [
    "fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Refactor using nn.Linear\n",
    "\n",
    "> **self.weights & self.bias & xb @ self.weights + self.bias --> nn.Linear**\n",
    "\n",
    "Instead of manually defining and initializing *self.weights* and *self.bias*, and calculating *xb @ self.weights + self.bias*, we will use the PyTorch class ***nn.Linear*** for a linear layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mnist_Logistic(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.linear = nn.Linear(784, 10)  # change\n",
    "        \n",
    "    def forward(self, xb):\n",
    "        return self.linear(xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(2.3254, grad_fn=<NllLossBackward>)\n"
     ]
    }
   ],
   "source": [
    "model = Mnist_Logistic()\n",
    "print(loss_func(model(xb), yb))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Refactor using optim\n",
    "\n",
    "> **param -= param.grad * lr & model.zero_grad() --> torch.optim**\n",
    "\n",
    "Instead of manually updating each parameter, we can use ***step*** method in *torch.optim* to take a forward step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(2.2828, grad_fn=<NllLossBackward>)\n"
     ]
    }
   ],
   "source": [
    "model = Mnist_Logistic()\n",
    "optimizer = optim.SGD(model.parameters(), lr=lr)\n",
    "print(loss_func(model(xb), yb))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0 --- i:    1 --- loss:  2.3240\n",
      "Epoch: 0 --- i:  301 --- loss:  1.9733\n",
      "Epoch: 0 --- i:  601 --- loss:  1.7667\n",
      "Epoch: 1 --- i:    1 --- loss:  1.7062\n",
      "Epoch: 1 --- i:  301 --- loss:  1.4404\n",
      "Epoch: 1 --- i:  601 --- loss:  1.3340\n",
      "Epoch: 2 --- i:    1 --- loss:  1.3534\n",
      "Epoch: 2 --- i:  301 --- loss:  1.1556\n",
      "Epoch: 2 --- i:  601 --- loss:  1.0871\n",
      "Epoch: 3 --- i:    1 --- loss:  1.1438\n",
      "Epoch: 3 --- i:  301 --- loss:  0.9905\n",
      "Epoch: 3 --- i:  601 --- loss:  0.9338\n",
      "Epoch: 4 --- i:    1 --- loss:  1.0081\n",
      "Epoch: 4 --- i:  301 --- loss:  0.8857\n",
      "Epoch: 4 --- i:  601 --- loss:  0.8308\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range((n - 1) // bs + 1):\n",
    "        xb, yb = x_train[i * bs: (i + 1) * bs], y_train[i * bs: (i + 1) * bs]\n",
    "        pred = model(xb)\n",
    "        loss = loss_func(pred, yb)\n",
    "\n",
    "        if i % 300 == 0:\n",
    "            print(f'Epoch: {epoch} --- i: {i + 1: 4d} --- loss: {loss.item(): .4f}')\n",
    "            \n",
    "        loss.backward()\n",
    "        optimizer.step()       # change\n",
    "        optimizer.zero_grad()  # change"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Refactor using Dataset\n",
    "\n",
    "> **x_train[i \\* bs, (i + 1) \\* bs] & y_train[...] --> torch.utils.data.TensorDataset**\n",
    "\n",
    "PyTorch's ***TensorDataset*** is a Dataset wrapping tensors. By defining a length and way of indexing, this also gives us a way to iterate, index, and slice along the 1st dimension of a tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = TensorDataset(x_train, y_train)  # TensorDataset VS Dataset： 前者继承后者，后者不宜直接使用，前者可直接使用于多个Tensor，后者不可"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Mnist_Logistic()\n",
    "optimizer = optim.SGD(model.parameters(), lr=lr)\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    for i in range((n - 1) // bs + 1):\n",
    "        xb, yb = train_ds[i * bs: (i + 1) * bs]  # Change: xb, yb一起从train_ds中indexing代替它们分别indexing\n",
    "        pred = model(xb)\n",
    "        loss = loss_func(pred, yb)\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Refactor using DataLoader\n",
    "\n",
    "> **for i in range((n - 1) // bs + 1) --> torch.utils.data.DataLoader**\n",
    "\n",
    "PyTorch's *DataLoader* is responsible for **managing batches**. We can create a *DataLoader* from any *Dataset*. *DataLoader* gives us each mini-batch automatically, making it easier to iterate over batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = TensorDataset(x_train, y_train)\n",
    "train_dl = DataLoader(train_ds, batch_size=bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Mnist_Logistic()\n",
    "optimizer = optim.SGD(model.parameters(), lr=lr)\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    for xb, yb in train_dl:  # change: 直接使用iterable的train_dl，代替手动计算循环次数和对应的indices！\n",
    "        pred = model(xb)\n",
    "        loss = loss_func(pred, yb)\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thanks to PyTorch's ***nn.Module, nn.Parameter, Dataset and DataLoader***, our training loop is now dramatically smaller and easier to understand."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Add Validation\n",
    "\n",
    "In training loop, we always should also have a **validation set**, in order to identify if it's overfitting.\n",
    "\n",
    "**Shuffling the training data** is important to prevent correlation between batches and overfitting. We don't shuffle the validation data.\n",
    "\n",
    "For the validation set, we'll use **a batch size that is twice as large as that for the training set**. This is because the validation set doesn't need backpropagation and thus takes less memory (it doesn't need to store the memory)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = TensorDataset(x_train, y_train)\n",
    "train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)  # shuffle=True\n",
    "\n",
    "valid_ds = TensorDataset(x_valid, y_valid)\n",
    "valid_dl = DataLoader(valid_ds, batch_size=bs * 2)            # shuffle=False, batch_size变为2倍"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:   0  val_loss:  1.6839\n",
      "Epoch:   1  val_loss:  1.3178\n",
      "Epoch:   2  val_loss:  1.0999\n",
      "Epoch:   3  val_loss:  0.9595\n",
      "Epoch:   4  val_loss:  0.8627\n"
     ]
    }
   ],
   "source": [
    "model = Mnist_Logistic()\n",
    "optimizer = optim.SGD(model.parameters(), lr=lr)\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    # Training\n",
    "    model.train()            # 之前没这句：默认是Training Phase；现在加这句：各epoch之间要来回切换\n",
    "    for xb, yb in train_dl:\n",
    "        pred = model(xb)\n",
    "        loss = loss_func(pred, yb)\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "    \n",
    "    # Validation\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        valid_loss = sum(loss_func(model(xb), yb) for xb, yb in valid_dl)\n",
    "    print(f'Epoch: {epoch: 3d}  val_loss: {valid_loss / len(valid_dl): .4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Create fit() and get_date()\n",
    "\n",
    "> 刘尧：尽可能把**重复且固定不变的操作**封闭成方法，这样在开发调试时可便捷地多次使用！如**生成DataLoader，计算loss，获得模型和optimizer，训练模型**等。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_batch(model, loss_func, xb, yb, optimizer=None):\n",
    "    loss = loss_func(model(xb), yb)\n",
    "    if optimizer is not None:\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "    return loss.item(), len(xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "def fit(epochs, model, loss_func, optimizer, train_dl, valid_dl):\n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        for xb, yb in train_dl:\n",
    "            loss_batch(model, loss_func, xb, yb, optimizer)\n",
    "            \n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            losses, nums = zip(*[loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl])\n",
    "        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)\n",
    "        print(f'Epoch: {epoch}  val_loss: {val_loss: .5f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataloader(train_ds, valid_ds, bs):\n",
    "    return (DataLoader(train_ds, batch_size=bs, shuffle=True), DataLoader(valid_ds, batch_size=bs * 2), )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model_optimizer():\n",
    "    model = Mnist_Logistic()\n",
    "    optimizer = optim.SGD(model.parameters(), lr=lr)\n",
    "    return model, optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, our whole process of obtaining the data loaders and fitting the model can be run in **3 lines of code**, and we can use these 3 lines of code to train a variety of models!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0  val_loss:  0.35634\n",
      "Epoch: 1  val_loss:  0.31089\n",
      "Epoch: 2  val_loss:  0.29377\n",
      "Epoch: 3  val_loss:  0.28763\n",
      "Epoch: 4  val_loss:  0.28284\n"
     ]
    }
   ],
   "source": [
    "train_dl, valid_dl = get_dataloader(train_ds, valid_ds, bs)\n",
    "model, optimizer = get_model_optimizer()\n",
    "fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Switch to CNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mnist_CNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1)\n",
    "        self.conv2 = nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1)\n",
    "        self.conv3 = nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1)\n",
    "        \n",
    "    def forward(self, xb):\n",
    "        xb = xb.view(-1, 1, 28, 28)\n",
    "        xb = F.relu(self.conv1(xb))\n",
    "        xb = F.relu(self.conv2(xb))\n",
    "        xb = F.relu(self.conv3(xb))\n",
    "        xb = F.avg_pool2d(xb, 4)\n",
    "        return xb.view(-1, xb.size(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0  val_loss:  0.43858\n",
      "Epoch: 1  val_loss:  0.32135\n",
      "Epoch: 2  val_loss:  0.22175\n",
      "Epoch: 3  val_loss:  0.19849\n",
      "Epoch: 4  val_loss:  0.18330\n"
     ]
    }
   ],
   "source": [
    "model = Mnist_CNN()\n",
    "lr = 0.1\n",
    "optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)\n",
    "fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)  # 之前定义的通用方法可多次使用，给开发调试带来很大便捷！"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 12. nn.Sequential\n",
    "\n",
    "torch.nn has another handy class we can use to simply our code: **Sequential**. To take advantage of this, we need to be able to easily define a **custom layer** from a given function. For instance, tensor.view.\n",
    "\n",
    "> 刘尧：Lambda类非常有用！可把Plain Python function转化为Layer/Module，从而加入成模型构建中！ PyTorch中没有自带的Lambda吗？？"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Lambda(nn.Module):\n",
    "    def __init__(self, func):\n",
    "        super().__init__()\n",
    "        self.func = func\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.func(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(\n",
    "    Lambda(lambda x: x.view(-1, 1, 28, 28)),\n",
    "    nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),\n",
    "    nn.ReLU(),\n",
    "    nn.AvgPool2d(4),\n",
    "    Lambda(lambda x: x.view(x.size(0), -1)),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0  val_loss:  0.37117\n",
      "Epoch: 1  val_loss:  0.22436\n",
      "Epoch: 2  val_loss:  0.20235\n",
      "Epoch: 3  val_loss:  0.16669\n",
      "Epoch: 4  val_loss:  0.16261\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)\n",
    "fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 13. Wrapping DataLoader\n",
    "\n",
    "Our CNN is fairly concise, but it only works with MNIST, because it assumes that **the input is a 28\\*28 long vector** and **the final CNN grid size is 4\\*4** (since that's the average pooling kernel size we used). Let's get rid of these 2 assumptions, so our model can work with any 2d single channel images. \n",
    "\n",
    "> 刘尧：把**动态的、可变化的部分**从模型构建中删除，且单独封闭成方法或类！比如输入数据的preprocessing。这种思想很重要，可让模型更加通用和灵活\n",
    "\n",
    "First, we can **remove the initial Lambda lyaer** but moving the data preprocessing into a generator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(x, y):                 # 此操作是动态的、可能会变化的，因为输入数据的size不唯一，不明确\n",
    "    return x.view(-1, 1, 28, 28), y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WrappedDataLoader:\n",
    "    def __init__(self, dl, func):  # 此处把动态的、可能会变化的外部操作传入！\n",
    "        self.dl = dl\n",
    "        self.func = func\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.dl)\n",
    "    \n",
    "    def __iter__(self):\n",
    "        batches = iter(self.dl)\n",
    "        for b in batches:\n",
    "            yield (self.func(*b))  # 注意yield的使用！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl, valid_dl = get_dataloader(train_ds, valid_ds, bs)\n",
    "train_dl = WrappedDataLoader(train_dl, preprocess)\n",
    "valid_dl = WrappedDataLoader(valid_dl, preprocess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we can replace ***nn.AvgPool2d*** with ***nn.AdaptiveAvgPool2d***, which allows us to define the size of the output tensor we want, as a result, our model will work with any size input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(\n",
    "    nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),\n",
    "    nn.ReLU(),\n",
    "    nn.AdaptiveAvgPool2d(1),\n",
    "    Lambda(lambda x: x.view(x.size(0), -1)),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0  val_loss:  0.67185\n",
      "Epoch: 1  val_loss:  0.56409\n",
      "Epoch: 2  val_loss:  0.52121\n",
      "Epoch: 3  val_loss:  0.48000\n",
      "Epoch: 4  val_loss:  0.46667\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)\n",
    "fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 14. Using your GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else torch.device('cpu'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(x, y):\n",
    "    return x.view(-1, 1, 28, 28).to(device), y.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl, valid_dl = get_dataloader(train_ds, valid_ds, bs)\n",
    "train_dl = WrappedDataLoader(train_dl, preprocess)\n",
    "valid_dl = WrappedDataLoader(valid_dl, preprocess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0  val_loss:  0.46358\n",
      "Epoch: 1  val_loss:  0.43656\n",
      "Epoch: 2  val_loss:  0.43888\n",
      "Epoch: 3  val_loss:  0.42316\n",
      "Epoch: 4  val_loss:  0.43026\n"
     ]
    }
   ],
   "source": [
    "model.to(device)\n",
    "optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)\n",
    "fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 15. Closing thoughts\n",
    "\n",
    "We now have a general data pipeline and training loop which we can use for training many types of models using PyTorch.\n",
    "\n",
    "Of course, there are many things that can be added, such as **data augmentation, hyperparameter tuning, monitoring training, transfer learning, and so forth**. These features are available in the ***fastai*** library, which has been developed using the same design approach shown in this tutorial, providing a natural next step for practitioners.\n",
    "\n",
    "#### Summary\n",
    "\n",
    "- **torch.nn.Module**: creates a callable which behaves like a function, but can also **contain state** (such as neural net layer weights). It knows what parameter(s) it contains and can zero all their gradient, loop through them for weight updates, etc.\n",
    "\n",
    "- torch.nn.Parameter: a wrapper for a tensor that **tells a Module** that it has weights that need updating during backprop.\n",
    "\n",
    "- torch.nn.functional: a module which contains activation function, loss functions, etc, as well as **non-stateful versions of layers** such as convolutional and linear layers.\n",
    "\n",
    "> 刘尧：functional中另外version的卷积、线性和池化等操作，如F.conv1d, F.max_pool2d, F.linear等，是**non-stateful！ 这是嘛意思？不需要weights，更不需要进行updating？**\n",
    "\n",
    "- torch.optim: contains optimizers such as SGD, which update the weights of Parameter during teh backward step.\n",
    "\n",
    "- torch.utils.data.Dataset: an abstract interface of objects with a \\_\\_len\\_\\_ and a \\_\\_getitem\\_\\_, including classes provided with PyTorch such as ***TensorDataset***.\n",
    "\n",
    "- torch.utils.data.DataLoader: takes any Dataset and creates an iterator which returns batches of data."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
